Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add redundant broadcast elimination pass after stablehlo conversion. #1252

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add redundant broadcast elimination pass after stablehlo conversion. #1252

wants to merge 1 commit into from
+107 −0

Conversation

uazizTT
Copy link
Contributor

@uazizTT uazizTT commented Nov 13, 2024 •
edited
Loading

Stablehlo can introduce broadcast ops for cases that do not require shape conversion. This can introduce an error during conversion as multiple operands of the consumer might need to be folded. This pass removes the redundant broadcasts that prevents an error later.

Fixes #1235

Example:

#any_device_tile = #tt.operand_constraint<dram|l1|tile|none|interleaved|single_bank|height_sharded|width_sharded|block_sharded|any_layout|any_device_tile>
module {
  func.func @main(%arg0: tensor<8xf32>, %arg1: tensor<1xf32>) -> tensor<8xf32> {
    %0 = tensor.empty() : tensor<8xf32>
    %1 = "ttir.broadcast"(%arg0, %0) <{dimension = [0], operand_constraints = [#any_device_tile, #any_device_tile]}> : (tensor<8xf32>, tensor<8xf32>) -> tensor<8xf32>
    %2 = tensor.empty() : tensor<8xf32>
    %3 = "ttir.broadcast"(%arg1, %2) <{dimension = [0], operand_constraints = [#any_device_tile, #any_device_tile]}> : (tensor<1xf32>, tensor<8xf32>) -> tensor<8xf32>
    %4 = tensor.empty() : tensor<8xf32>
    %5 = "ttir.add"(%1, %3, %4) <{operandSegmentSizes = array<i32: 2, 1>, operand_constraints = [#any_device_tile, #any_device_tile, #any_device_tile]}> : (tensor<8xf32>, tensor<8xf32>, tensor<8xf32>) -> tensor<8xf32>
    return %5 : tensor<8xf32>
  }
}

After this pass:

// -----// IR Dump After RedundantBroadcastElimination (redundant-broadcast-elimination) //----- //
#any_device_tile = #tt.operand_constraint<dram|l1|tile|none|interleaved|single_bank|height_sharded|width_sharded|block_sharded|any_layout|any_device_tile>
module {
  func.func @main(%arg0: tensor<8xf32>, %arg1: tensor<1xf32>) -> tensor<8xf32> {
    %0 = tensor.empty() : tensor<8xf32>
    %1 = tensor.empty() : tensor<8xf32>
    %2 = "ttir.broadcast"(%arg1, %1) <{dimension = [0], operand_constraints = [#any_device_tile, #any_device_tile]}> : (tensor<1xf32>, tensor<8xf32>) -> tensor<8xf32>
    %3 = tensor.empty() : tensor<8xf32>
    %4 = "ttir.add"(%arg0, %2, %3) <{operandSegmentSizes = array<i32: 2, 1>, operand_constraints = [#any_device_tile, #any_device_tile, #any_device_tile]}> : (tensor<8xf32>, tensor<8xf32>, tensor<8xf32>) -> tensor<8xf32>
    return %4 : tensor<8xf32>
  }
}

@uazizTT uazizTT added the stablehlo conversion bug Bugs in StableHLO conversion label Nov 13, 2024
@uazizTT uazizTT self-assigned this Nov 13, 2024
nsmithtt
nsmithtt reviewed Nov 14, 2024
View reviewed changes
lib/Conversion/StableHLOToTTIR/RedundantBroadcastElimination.cpp
ModuleOp module = getOperation();
IRRewriter rewriter(&getContext());

module->walk([&](Operation *op) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think you can pattern match directly w/ walk by writing:

module->walk([&](mlir::tt::ttir::BroadcastOp op) {

lib/Conversion/StableHLOToTTIR/RedundantBroadcastElimination.cpp
return;
}

if (op->getResult(0).getType() == op->getOperand(0).getType()) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hmm, this was not what I expected the name RedundantBroadcastEliminationPass to mean. We need a simple test demonstrating this case.

To me this implies that the broadcast is operating on a tensor that's already broadcasted in that dimension, but should this be illegal by construction?

mmanzoorTT
mmanzoorTT reviewed Nov 15, 2024
View reviewed changes
include/ttmlir/Conversion/Passes.h
@@ -7,6 +7,7 @@

#ifdef TTMLIR_ENABLE_STABLEHLO
#include "ttmlir/Conversion/ArithToStableHLO/ArithToStableHLO.h"
#include "ttmlir/Conversion/RedundantBroadcastElimination/RedundantBroadcastElimination.h"
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think, we should make it a generic pass instead of StableHLO dependent pass. Other dialects may also require this optimization.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@mmanzoorTT mmanzoorTT mmanzoorTT left review comments

@nsmithtt nsmithtt nsmithtt left review comments

@AleksKnezevic AleksKnezevic Awaiting requested review from AleksKnezevic

@sdjordjevicTT sdjordjevicTT Awaiting requested review from sdjordjevicTT sdjordjevicTT is a code owner

@rpavlovicTT rpavlovicTT Awaiting requested review from rpavlovicTT rpavlovicTT is a code owner

@svuckovicTT svuckovicTT Awaiting requested review from svuckovicTT svuckovicTT is a code owner

@mtopalovicTT mtopalovicTT Awaiting requested review from mtopalovicTT mtopalovicTT is a code owner

@mrakitaTT mrakitaTT Awaiting requested review from mrakitaTT mrakitaTT is a code owner

@nobradovictt nobradovictt Awaiting requested review from nobradovictt nobradovictt is a code owner

@jserbedzijaTT jserbedzijaTT Awaiting requested review from jserbedzijaTT jserbedzijaTT is a code owner

At least 1 approving review is required to merge this pull request.

Assignees

@uazizTT uazizTT

Labels
stablehlo conversion bug Bugs in StableHLO conversion
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

Stablehlo Assertion `!impl->wasOpReplaced(op) && "attempting to modify a replaced/erased op"' failed.
3 participants
@uazizTT @nsmithtt @mmanzoorTT